import torch
import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 3)
        self.conv2 = nn.Conv2d(16, 32, 3)
        self.conv3 = nn.Conv2d(32, 64, 3)
        self.conv4 = nn.Conv2d(64, 128, 3)
        self.conv5 = nn.Conv2d(128, 256, 3)
        self.fc = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv3(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.relu(self.conv5(x))
        x = F.max_pool2d(x, 2, 2)
        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(x.size(0), -1)
        logits = self.fc(x)
        output = F.softmax(logits, dim=1)
        return output

class DeCNN(nn.Module):
    def __init__(self):
        super(DeCNN, self).__init__()
        self.deconv1 = nn.ConvTranspose2d(256, 128, 3)
        self.deconv2 = nn.ConvTranspose2d(128, 64, 3)
        self.deconv3 = nn.ConvTranspose2d(64, 32, 3)
        self.deconv4 = nn.ConvTranspose2d(32, 16, 3)
        self.deconv5 = nn.ConvTranspose2d(16, 3, 3)

    def forward(self, x):
        x = F.relu(self.deconv1(x))
        x = F.interpolate(x, scale_factor=2)
        x = F.relu(self.deconv2(x))
        x = F.interpolate(x, scale_factor=2)
        x = F.relu(self.deconv3(x))
        x = F.interpolate(x, scale_factor=2)
        x = F.relu(self.deconv4(x))
        x = F.interpolate(x, scale_factor=2)
        x = F.relu(self.deconv5(x))
        return x

# 实例化图像分类网络和重构网络
net = CNN()
decnn = DeCNN()

# 定义交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 示例输入数据
input_data = torch.randn(1, 3, 250, 250)  # 假设输入数据为3通道，大小为32x32

# 示例目标标签
target = torch.LongTensor([3])  # 假设目标类别为3

# 前向传播
output = net(input_data)

# 计算损失
loss = criterion(output, target)

# 重构网络前向传播
reconstructed_data = decnn(output)

print(loss)
